refactor(gemm): split the gemv_q8 grab-bag by family - #24
Merged
Conversation
Contributor
|
@ekryski clean split, no kernel dropped, no bench-name drift. The body understates one thing: the 21 "not deleted, registry-bound" kernels were renamed |
TheTom
force-pushed
the
ek/kernel-reorg-gemm-quantized
branch
from
June 22, 2026 21:27
2b24cef to
742dcea
Compare
gemv_q8.rs bundled 25 kernels spanning 6 families. Split each to its home:
- gemm (12 + 4 benches) -> kernels/gemm/gemv_quantized.rs: the Q8/Q4
inline-dequant GEMVs (plain/coalesced/relu2/accum, grouped(_rows,_tiled),
Q4 _vec/_2row)
- moe (5) -> NEW kernels/moe/: gather_q4.rs (batched expert up/down/
weighted-sum) + sigmoid_bias.rs (router pre-score). moe folder seeded
early; the rest of the moe family migrates in wave 3.
- ops (4) -> kernels/ops/: cast_f32_f16 / cast_f16_f32 / vscale (unary.rs),
slice (copy.rs)
- ssm (2) -> kernels/ssm/scan.rs: softplus_add, gated_group_rmsnorm
(next to their _rows / _batched siblings)
- conv (1) -> kernels/convolution/conv1d_causal.rs: conv_roll (state roll)
- kv (1) -> kernels/kv_cache/cache.rs: kv_append
All ffai_ -> mt_. No kernels deleted (the 21 consumer-less kernels are
runtime-dispatch/registry-bound). Fixed the 3 gemv_q8 correctness-test imports.
Dissolves the last parked grab-bag from the gemm migration.
TheTom
force-pushed
the
ek/kernel-reorg-gemv-q8-split
branch
from
June 22, 2026 21:27
3255164 to
b86ce1e
Compare
TheTom
pushed a commit
that referenced
this pull request
Jun 22, 2026
Move all 35 moe-family files from ffai/ + mlx/ into kernels/moe/ (which #24 seeded with gather_q4 + sigmoid_bias): - orchestration (ex moe.rs): router_topk + permute/unpermute + 10 gather_qmm - routers: router_topk_biased (ex dsv4_router_topk), router_sigmoid_bias, router_sqrtsoftplus, sigmoid_bias - mpp grouped BGEMM: mpp(+int8/bm8/bm64/×int8/×block_scaled) + mpp_shared - gguf-format expert matmul: bgemm_{q2k,iq2xxs,q4}_*, gemv_{rows,ws}_*, gather_* - down combine: down_swiglu_accum, down_weighted_sum_f16 - expert-indexed + block-scaled: dequant_gemv_expert_indexed(_block_scaled), block_scaled (ex mlx/block_scaled_moe) Filenames drop the redundant moe_ prefix (folder provides it); kernel names keep mt_moe_*. Model-name purge: mt_dsv4_router_topk -> mt_moe_router_topk_biased (distinct from the generic mt_moe_router_topk: selects by the biased score, weights by the unbiased). Bare dequant_gemv_int4_expert_indexed -> mt_ prefix. Fixed mpp_shared intra-imports and ~27 consumer test files (grouped/mixed use-blocks included). Format-axis fold (§7) deferred. orchestration.rs (~4k lines) moves whole here; split follows next.
TheTom
pushed a commit
that referenced
this pull request
Jun 22, 2026
Move all 35 moe-family files from ffai/ + mlx/ into kernels/moe/ (which #24 seeded with gather_q4 + sigmoid_bias): - orchestration (ex moe.rs): router_topk + permute/unpermute + 10 gather_qmm - routers: router_topk_biased (ex dsv4_router_topk), router_sigmoid_bias, router_sqrtsoftplus, sigmoid_bias - mpp grouped BGEMM: mpp(+int8/bm8/bm64/×int8/×block_scaled) + mpp_shared - gguf-format expert matmul: bgemm_{q2k,iq2xxs,q4}_*, gemv_{rows,ws}_*, gather_* - down combine: down_swiglu_accum, down_weighted_sum_f16 - expert-indexed + block-scaled: dequant_gemv_expert_indexed(_block_scaled), block_scaled (ex mlx/block_scaled_moe) Filenames drop the redundant moe_ prefix (folder provides it); kernel names keep mt_moe_*. Model-name purge: mt_dsv4_router_topk -> mt_moe_router_topk_biased (distinct from the generic mt_moe_router_topk: selects by the biased score, weights by the unbiased). Bare dequant_gemv_int4_expert_indexed -> mt_ prefix. Fixed mpp_shared intra-imports and ~27 consumer test files (grouped/mixed use-blocks included). Format-axis fold (§7) deferred. orchestration.rs (~4k lines) moves whole here; split follows next.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #23. Dissolves the last parked grab-bag from the gemm migration.
ffai/gemv_q8.rsbundled 25 kernels spanning 6 families in one 1069-line file. This splits each kernel to its proper family home.kernels/gemm/gemv_quantized.rs(12 + 4 benches)_coalesced/_relu2/_accum,grouped(_rows,_tiled), Q4_vec/_2rowkernels/moe/(5)gather_q4.rs(batched expert up/down/weighted-sum) +sigmoid_bias.rs(router pre-score)kernels/ops/(4)cast_f32_f16/cast_f16_f32/vscale(unary.rs),slice(copy.rs)kernels/ssm/scan.rs(2)softplus_add,gated_group_rmsnorm(next to_rows/_batchedsiblings)kernels/convolution/conv1d_causal.rs(1)conv_roll(causal-conv state roll)kernels/kv_cache/cache.rs(1)kv_appendNotes
ffai_→mt_. No kernels deleted — only the 4 q8-gemv kernels have Rust consumers; the other 21 are runtime-dispatch/registry-bound (deleting them would break Swift-side dispatch by name).moe/is seeded early (folder created ahead of the wave-3 moe migration); the remaining ~30moe_*files inffai/migrate then.moe_sigmoid_biasis not a dup of the existingmoe_router_sigmoid_bias(different output shape).Verification
cargo build --workspace --bins✓cargo clippy --workspace --all-targets --all-features -- -D warnings✓ (clean)tile build— all 1272 kernels codegen to MSL ✓; every relocated kernel registered under itsmt_name (tile inspect)